"""
JWT Handler
===========

JWT token management utilities for authentication and authorization.
Provides token generation, validation, and refresh functionality.
"""

import jwt
import logging
from datetime import datetime, timedelta
from typing import Optional, Dict, Any, Union
from flask import current_app, request, jsonify
from functools import wraps

# Initialize logger
logger = logging.getLogger(__name__)

class JWTHandler:
    """
    JWT token handler for authentication and authorization.
    
    Features:
    - Token generation and validation
    - Refresh token management
    - Blacklist functionality
    - Role-based access control
    """
    
    def __init__(self):
        """Initialize JWT handler."""
        self.secret_key = current_app.config.get('JWT_SECRET_KEY', 'default-secret-key')
        self.algorithm = 'HS256'
        self.access_token_expiry = current_app.config.get('JWT_ACCESS_TOKEN_EXPIRES', timedelta(hours=24))
        self.refresh_token_expiry = current_app.config.get('JWT_REFRESH_TOKEN_EXPIRES', timedelta(days=30))
    
    def generate_tokens(self, user_id: int, username: str, email: str, 
                       roles: Optional[list] = None, permissions: Optional[list] = None) -> Dict[str, str]:
        """
        Generate access and refresh tokens.
        
        Args:
            user_id: User ID
            username: Username
            email: User email
            roles: User roles (optional)
            permissions: User permissions (optional)
        
        Returns:
            Dict with access_token and refresh_token
        """
        now = datetime.utcnow()
        
        # Access token payload
        access_payload = {
            'user_id': user_id,
            'username': username,
            'email': email,
            'roles': roles or [],
            'permissions': permissions or [],
            'type': 'access',
            'iat': now,
            'exp': now + self.access_token_expiry
        }
        
        # Refresh token payload
        refresh_payload = {
            'user_id': user_id,
            'username': username,
            'type': 'refresh',
            'iat': now,
            'exp': now + self.refresh_token_expiry
        }
        
        # Generate tokens
        access_token = jwt.encode(access_payload, self.secret_key, algorithm=self.algorithm)
        refresh_token = jwt.encode(refresh_payload, self.secret_key, algorithm=self.algorithm)
        
        logger.info(f"Generated tokens for user {username}")
        
        return {
            'access_token': access_token,
            'refresh_token': refresh_token,
            'token_type': 'Bearer',
            'expires_in': int(self.access_token_expiry.total_seconds())
        }
    
    def verify_token(self, token: str, token_type: str = 'access') -> Optional[Dict[str, Any]]:
        """
        Verify and decode JWT token.
        
        Args:
            token: JWT token to verify
            token_type: Expected token type ('access' or 'refresh')
        
        Returns:
            Decoded payload or None if invalid
        """
        try:
            payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm])
            
            # Check token type
            if payload.get('type') != token_type:
                logger.warning(f"Invalid token type. Expected {token_type}, got {payload.get('type')}")
                return None
            
            # Check if token is blacklisted
            if self.is_token_blacklisted(token):
                logger.warning("Token is blacklisted")
                return None
            
            return payload
        
        except jwt.ExpiredSignatureError:
            logger.warning("Token has expired")
            return None
        except jwt.InvalidTokenError as e:
            logger.warning(f"Invalid token: {str(e)}")
            return None
        except Exception as e:
            logger.error(f"Error verifying token: {str(e)}")
            return None
    
    def refresh_access_token(self, refresh_token: str) -> Optional[Dict[str, str]]:
        """
        Generate new access token using refresh token.
        
        Args:
            refresh_token: Valid refresh token
        
        Returns:
            Dict with new access_token or None if invalid
        """
        payload = self.verify_token(refresh_token, token_type='refresh')
        
        if not payload:
            return None
        
        # Generate new access token
        now = datetime.utcnow()
        access_payload = {
            'user_id': payload['user_id'],
            'username': payload['username'],
            'type': 'access',
            'iat': now,
            'exp': now + self.access_token_expiry
        }
        
        access_token = jwt.encode(access_payload, self.secret_key, algorithm=self.algorithm)
        
        logger.info(f"Refreshed access token for user {payload['username']}")
        
        return {
            'access_token': access_token,
            'token_type': 'Bearer',
            'expires_in': int(self.access_token_expiry.total_seconds())
        }
    
    def blacklist_token(self, token: str) -> bool:
        """
        Add token to blacklist.
        
        Args:
            token: Token to blacklist
        
        Returns:
            bool: True if successful
        """
        try:
            # Decode token to get expiry
            payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm])
            exp = payload.get('exp')
            
            if exp:
                # Calculate TTL
                ttl = exp - int(datetime.utcnow().timestamp())
                if ttl > 0:
                    # Add to Redis blacklist with TTL
                    from app.utils.redis_client import RedisClient
                    from flask import current_app
                    
                    redis_client = RedisClient(
                        host=current_app.config.get('REDIS_HOST', 'localhost'),
                        port=current_app.config.get('REDIS_PORT', 6379),
                        db=4  # Use separate DB for blacklisted tokens
                    )
                    
                    redis_client.set(f"blacklist:{token}", "1", ex=ttl)
                    logger.info(f"Blacklisted token with TTL {ttl}")
                    return True
            
            return False
        
        except Exception as e:
            logger.error(f"Error blacklisting token: {str(e)}")
            return False
    
    def is_token_blacklisted(self, token: str) -> bool:
        """
        Check if token is blacklisted.
        
        Args:
            token: Token to check
        
        Returns:
            bool: True if blacklisted
        """
        try:
            from app.utils.redis_client import RedisClient
            from flask import current_app
            
            redis_client = RedisClient(
                host=current_app.config.get('REDIS_HOST', 'localhost'),
                port=current_app.config.get('REDIS_PORT', 6379),
                db=4
            )
            
            return redis_client.exists(f"blacklist:{token}") > 0
        
        except Exception as e:
            logger.error(f"Error checking token blacklist: {str(e)}")
            return False
    
    def extract_token_from_header(self, auth_header: str) -> Optional[str]:
        """
        Extract JWT token from Authorization header.
        
        Args:
            auth_header: Authorization header value
        
        Returns:
            Token string or None if invalid format
        """
        if not auth_header:
            return None
        
        try:
            parts = auth_header.split()
            if len(parts) == 2 and parts[0].lower() == 'bearer':
                return parts[1]
        except Exception:
            pass
        
        return None

# Global JWT handler instance
jwt_handler = JWTHandler()

def login_required(f):
    """
    Decorator to require authentication for a route.
    
    Args:
        f: Function to decorate
    
    Returns:
        Decorated function
    """
    @wraps(f)
    def decorated_function(*args, **kwargs):
        auth_header = request.headers.get('Authorization')
        
        if not auth_header:
            return jsonify({
                'status': 'error',
                'message': 'Authorization header is required',
                'error_code': 'missing_auth_header'
            }), 401
        
        token = jwt_handler.extract_token_from_header(auth_header)
        
        if not token:
            return jsonify({
                'status': 'error',
                'message': 'Invalid authorization header format',
                'error_code': 'invalid_auth_format'
            }), 401
        
        payload = jwt_handler.verify_token(token)
        
        if not payload:
            return jsonify({
                'status': 'error',
                'message': 'Invalid or expired token',
                'error_code': 'invalid_token'
            }), 401
        
        # Add user info to request context
        request.current_user = payload
        
        return f(*args, **kwargs)
    
    return decorated_function

def roles_required(*required_roles):
    """
    Decorator to require specific roles for a route.
    
    Args:
        *required_roles: Required role names
    
    Returns:
        Decorated function
    """
    def decorator(f):
        @wraps(f)
        @login_required
        def decorated_function(*args, **kwargs):
            user_roles = request.current_user.get('roles', [])
            
            if not any(role in user_roles for role in required_roles):
                return jsonify({
                    'status': 'error',
                    'message': 'Insufficient permissions',
                    'error_code': 'insufficient_permissions',
                    'required_roles': list(required_roles),
                    'user_roles': user_roles
                }), 403
            
            return f(*args, **kwargs)
        
        return decorated_function
    
    return decorator

def permissions_required(*required_permissions):
    """
    Decorator to require specific permissions for a route.
    
    Args:
        *required_permissions: Required permission names
    
    Returns:
        Decorated function
    """
    def decorator(f):
        @wraps(f)
        @login_required
        def decorated_function(*args, **kwargs):
            user_permissions = request.current_user.get('permissions', [])
            
            if not any(permission in user_permissions for permission in required_permissions):
                return jsonify({
                    'status': 'error',
                    'message': 'Insufficient permissions',
                    'error_code': 'insufficient_permissions',
                    'required_permissions': list(required_permissions),
                    'user_permissions': user_permissions
                }), 403
            
            return f(*args, **kwargs)
        
        return decorated_function
    
    return decorator